-
Notifications
You must be signed in to change notification settings - Fork 7.9k
Fix #7100: Add cancellation check in run_stream() to stop iteration i… #7123
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Fix #7100: Add cancellation check in run_stream() to stop iteration i… #7123
Conversation
…eration immediately - Add cancellation token check before yielding messages in run_stream() - Prevents processing buffered messages after cancellation token is cancelled - Fixes issue where signal handlers (SIGINT/Ctrl+C) couldn't stop iteration immediately - Add test_round_robin_group_chat_run_stream_cancellation to verify fix The issue was that run_stream() would continue processing all buffered messages from the queue even after the cancellation token was cancelled, because the cancellation check only happened when awaiting the next message. This fix adds a check after receiving each message to break immediately if cancellation occurred.
|
@Shizoqua please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
Why are these changes needed?
This PR fixes issue #7100 where
run_stream()would continue processing all buffered messages from the queue even after theCancellationTokenwas cancelled, particularly when cancelled by signal handlers (e.g., SIGINT/Ctrl+C).Problem:
When a
CancellationTokenis cancelled (especially by signal handlers), therun_stream()async iteration loop continues processing all buffered messages before checking the cancellation status. This makes applications appear unresponsive to user interruption, as the loop drains the entire message queue before stopping.Root Cause:
In
_base_group_chat.py, therun_stream()method has awhile True:loop that only checks for cancellation when awaiting the next message from the queue. If there are buffered messages in_output_message_queue, they get processed and yielded even after the cancellation token is cancelled, because the cancellation check happens too late in the iteration cycle.Solution:
Added a cancellation check immediately after receiving each message from the queue, before processing or yielding it. This ensures that if the cancellation token was cancelled, the loop breaks immediately without processing any remaining buffered messages.
Code Change:
This fix ensures that when a signal handler cancels the token, the iteration stops immediately on the next loop iteration, providing the expected responsive behavior for interactive CLI applications.
Related issue number
Fixes #7100
Related to:
runandrun_stream, and add example to show how to use it #4029 - Respect Cancellation Token in AgentChat run and run_streamChecks
test_round_robin_group_chat_run_stream_cancellation()test intest_group_chat.pytest_round_robin_group_chat_cancellationstill passes (backward compatibility verified)poe format)poe lint)pytest packages/autogen-agentchat/tests/test_group_chat.py::test_round_robin_group_chat_run_stream_cancellation)Testing
New Test Added:
test_round_robin_group_chat_run_stream_cancellation()- Verifies thatrun_stream()stops immediately whenCancellationTokenis cancelled, preventing processing of buffered messages.Test Results:
single_threadedandembeddedruntime variants)run_streamtests continue to passFiles Changed
python/packages/autogen-agentchat/src/autogen_agentchat/teams/_group_chat/_base_group_chat.pyrun_stream()method (lines 550-554)python/packages/autogen-agentchat/tests/test_group_chat.pytest_round_robin_group_chat_run_stream_cancellation()test (lines 766-818)